package com.example.springaiollama.service;

import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Service
public class AiService {

    private final ChatClient chatClient;
    private final FunctionCallingService functionCallingService;

    @Autowired
    public AiService(ChatClient chatClient, FunctionCallingService functionCallingService) {
        this.chatClient = chatClient;
        this.functionCallingService = functionCallingService;
    }

    public String generateResponse(String userMessage) {
        Message message = new UserMessage(userMessage);
        Prompt prompt = new Prompt(List.of(message));
        
        ChatResponse response = chatClient.call(prompt);
        return response.getResult().getOutput().getContent();
    }

    public String generateResponseWithSystemPrompt(String userMessage, String systemPrompt) {
        Message systemMessage = new SystemMessage(systemPrompt);
        Message userMsg = new UserMessage(userMessage);
        Prompt prompt = new Prompt(List.of(systemMessage, userMsg));
        
        ChatResponse response = chatClient.call(prompt);
        return response.getResult().getOutput().getContent();
    }

    /**
     * 使用 Function Calling 生成响应
     * AI 可以根据需要使用注册的工具
     */
    public String generateResponseWithFunctionCalling(String userMessage) {
        // 构建包含工具信息的系统提示词
        String systemPrompt = buildSystemPromptWithTools();
        String aiResponse = generateResponseWithSystemPrompt(userMessage, systemPrompt);
        
        // 处理 AI 响应中的函数调用
        return processFunctionCalls(aiResponse);
    }

    /**
     * 使用 Function Calling 和系统提示词生成响应
     */
    public String generateResponseWithFunctionCallingAndSystemPrompt(String userMessage, String systemPrompt) {
        String enhancedSystemPrompt = systemPrompt + "\n\n" + buildSystemPromptWithTools();
        String aiResponse = generateResponseWithSystemPrompt(userMessage, enhancedSystemPrompt);
        
        // 处理 AI 响应中的函数调用
        return processFunctionCalls(aiResponse);
    }

    /**
     * 构建包含工具信息的系统提示词
     */
    private String buildSystemPromptWithTools() {
        Map<String, String> tools = functionCallingService.getAvailableTools();
        StringBuilder prompt = new StringBuilder();
        prompt.append("你可以使用以下工具来帮助回答问题：\n");
        
        for (Map.Entry<String, String> tool : tools.entrySet()) {
            prompt.append("- ").append(tool.getKey()).append(": ").append(tool.getValue()).append("\n");
        }
        
        prompt.append("\n如果用户需要这些功能，请直接调用相应的工具。例如：\n");
        prompt.append("- 如果用户问时间，调用 getCurrentTime()\n");
        prompt.append("- 如果用户问天气，调用 getWeather(城市名)\n");
        prompt.append("- 如果用户需要计算，调用 calculate(表达式)\n");
        prompt.append("- 如果用户需要翻译，调用 translate(文本, 目标语言)\n");
        
        return prompt.toString();
    }

    /**
     * 处理 AI 响应中的函数调用
     */
    private String processFunctionCalls(String aiResponse) {
        // 检查是否包含函数调用指示
        if (aiResponse.contains("getCurrentTime()") || aiResponse.contains("getCurrentTime")) {
            return "当前时间: " + functionCallingService.getCurrentTime();
        }
        
        if (aiResponse.contains("getWeather(") || aiResponse.contains("getWeather")) {
            // 提取城市名
            Pattern pattern = Pattern.compile("getWeather\\(([^)]+)\\)");
            Matcher matcher = pattern.matcher(aiResponse);
            if (matcher.find()) {
                String city = matcher.group(1).replaceAll("[\"']", "");
                return "天气信息: " + functionCallingService.getWeather(city);
            }
            // 如果没有找到括号，尝试从文本中提取城市名
            if (aiResponse.contains("北京")) return "天气信息: " + functionCallingService.getWeather("北京");
            if (aiResponse.contains("上海")) return "天气信息: " + functionCallingService.getWeather("上海");
            if (aiResponse.contains("广州")) return "天气信息: " + functionCallingService.getWeather("广州");
            if (aiResponse.contains("深圳")) return "天气信息: " + functionCallingService.getWeather("深圳");
        }
        
        if (aiResponse.contains("calculate(") || aiResponse.contains("calculate")) {
            // 提取表达式
            Pattern pattern = Pattern.compile("calculate\\(([^)]+)\\)");
            Matcher matcher = pattern.matcher(aiResponse);
            if (matcher.find()) {
                String expression = matcher.group(1).replaceAll("[\"']", "");
                return "计算结果: " + functionCallingService.calculate(expression);
            }
            // 如果没有找到括号，尝试从文本中提取数字和运算符
            Pattern calcPattern = Pattern.compile("(\\d+\\s*[+\\-*/]\\s*\\d+)");
            Matcher calcMatcher = calcPattern.matcher(aiResponse);
            if (calcMatcher.find()) {
                String expression = calcMatcher.group(1);
                return "计算结果: " + functionCallingService.calculate(expression);
            }
        }
        
        if (aiResponse.contains("translate(") || aiResponse.contains("translate")) {
            // 提取翻译参数
            Pattern pattern = Pattern.compile("translate\\(([^,]+),\\s*([^)]+)\\)");
            Matcher matcher = pattern.matcher(aiResponse);
            if (matcher.find()) {
                String text = matcher.group(1).replaceAll("[\"']", "");
                String targetLanguage = matcher.group(2).replaceAll("[\"']", "");
                return "翻译结果: " + functionCallingService.translate(text, targetLanguage);
            }
        }
        
        // 如果没有检测到函数调用，返回原始响应
        return aiResponse;
    }
} 